# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import copy
import math

from coin.exchange.base.kr_rest.currency import NativeCurrency
from coin.exchange.base.kr_rest.futures_product import NativeFuturesProduct
from coin.exchange.base.kr_rest.product import NativeProduct


class OrderGatewayInfo(object):
  CurrencyType = None
  ProductType = None
  BalanceType = float
  PositionType = float
  BalanceEmptyValue = 0.
  PositionEmptyValue = 0.

  def __init__(
      self,
      currencies,  # : List[CurrencyType],
      products):  # : List[ProductType]):
    # assert all([isinstance(x, self.CurrencyType) for x in currencies]), \
    #   (self.CurrencyType, [type(x) for x in currencies])
    # assert all([isinstance(x, self.ProductType) for x in products]), \
    #   (self.ProductType, [type(x) for x in products])
    assert len(currencies) > 0
    assert len(products) > 0
    # TODO: assert all currencies are in products

    self._currencies = currencies
    self._products = products
    self._balances = {currency.symbol: None for currency in currencies}
    self._positions = {product.symbol: None for product in products}
    self._balances_internal = {currency.symbol: None for currency in currencies}
    self._positions_internal = {product.symbol: None for product in products}

  def __eq__(self, other):
    if self is other:
      return True

    if not isinstance(other, type(self)):
      return False

    if set(self.products) != set(other.products):
      return False

    if set(self.currencies) != set(other.currencies):
      return False

    balance_same = []
    for currency in self.currencies:
      bal1 = self.get_balance(currency)
      bal2 = other.get_balance(currency)
      if bal1 is not None and bal2 is not None:
        balance_same.append(math.isclose(bal1, bal2))
      else:
        return False

    position_same = []
    for product in self.products:
      pos1 = self.get_position(product)
      pos2 = other.get_position(product)
      if pos1 is not None and pos2 is not None:
        position_same.append(math.isclose(pos1, pos2))
      else:
        return False

    return all(position_same) and all(balance_same)

  def assert_type_currency(self, currency):
    assert isinstance(currency, self.CurrencyType), (currency, type(currency))

  def assert_in_currency(self, currency):
    assert isinstance(currency, self.CurrencyType), (currency, type(currency))
    assert currency in self._currencies, (currency, self._currencies)

  def assert_type_product(self, product):
    assert isinstance(product, self.ProductType), (product, type(product))

  def assert_in_product(self, product):
    assert isinstance(product, self.ProductType), (product, type(product))
    assert product in self._products, (product, self._products)

  @property
  def currencies(self):
    return self._currencies

  @property
  def products(self):
    return self._products

  def get_currencies(self):
    return self._currencies

  def get_products(self):
    return self._products

  def get_balances(self):  # -> Dict[str, BalanceType]
    return self._balances

  def get_positions(self):  # -> Dict[str, PositionType]:
    return self._positions

  def get_balances_copy(self):  # -> Dict[str, BalanceType]
    return copy.copy(self._balances)

  def get_positions_copy(self):  # -> Dict[str, PositionType]:
    return copy.copy(self._positions)

  def get_notNone_balances_internal(self):
    return [x for x in self._balances_internal.values() if x is not None]

  def get_notNone_positions_internal(self):
    return [x for x in self._positions_internal.values() if x is not None]

  def get_balance(self, currency):
    # If type is right and not found, return empty value.
    # If type is str and not found, raise.
    if isinstance(currency, self.CurrencyType):
      return self._balances.get(currency.symbol, self.BalanceEmptyValue)
    elif isinstance(currency, NativeCurrency):
      # do not use get
      return self._balances[currency.symbol]
    elif isinstance(currency, str):
      # do not use get
      return self._balances[currency]
    else:
      raise TypeError("Not Supported: %s" % (currency,))

  def get_position(self, product):
    # If type is right and not found, return empty value.
    # If type is str and not found, raise.
    if isinstance(product, self.ProductType):
      return self._positions.get(product.symbol, self.PositionEmptyValue)
    elif isinstance(product, (NativeProduct, NativeFuturesProduct)):
      # do not use get
      return self._positions[product.symbol]
    elif isinstance(product, str):
      # do not use get
      return self._positions[product]
    else:
      raise TypeError("Not Supported: %s" % (product,))

  def balance_as_proto(self):
    return NotImplemented

  def position_as_proto(self):
    return NotImplemented


class BaseOrderGatewayInfoMutator(OrderGatewayInfo):
  def __init__(
      self,
      currencies,  # : List[CurrencyType],
      products):  # : List[ProductType]):
    super().__init__(currencies=currencies, products=products)

  def set_zero_balances(self):
    for currency in self._currencies:
      self.set_balance(currency, 0.)

  def reset_balance(self, currency):
    self.assert_type_currency(currency)
    self._balances[currency.symbol] = None
    self._balances_internal[currency.symbol] = None

  def set_balance(self, currency, balance, internal=None):
    self.assert_type_currency(currency)
    assert balance is None or isinstance(balance, self.BalanceType)
    self._balances[currency.symbol] = balance
    self._balances_internal[currency.symbol] = internal
    return balance

  def set_zero_positions(self):
    for product in self._products:
      self.set_position(product, 0.)

  def reset_position(self, product):
    self.assert_type_product(product)
    self._positions[product.symbol] = None
    self._positions_internal[product.symbol] = None

  def set_position(self, product, position, internal=None):
    self.assert_type_product(product)
    assert position is None or isinstance(position, self.PositionType)
    self._positions[product.symbol] = position
    self._positions_internal[product.symbol] = internal
    return position

  def reset(self):
    for currency in self.get_currencies():
      self.reset_balance(currency)
    for product in self.get_products():
      self.reset_position(product)
